import json
from typing import Dict
from pathlib import Path
from src.huffman import HuffmanTree

class HuffmanCompressor:
    def __init__(self, input_path: Path, output_path: Path):
        """哈夫曼树解压器初始化函数

        Args:
            input_path (Path): 输入文件
            output_path (Path): 输出文件
        """
        self.input_path = input_path
        self.output_path = output_path
        self.tree = HuffmanTree()
    
    def _get_frequency(self, content: str) -> Dict[str, int]:
        """根据文本内容生成字频字典

        Args:
            content (str): 文本内容

        Returns:
            Dict[str, int]: 生成的字频字典
        """
        freq_dict = {}
        for char in content:
            freq_dict[char] = freq_dict.get(char, 0) + 1
        return freq_dict
    
    def _to_byte_array(self, bit_string: str) -> bytes:
        """将二进制序列转换为字节序列

        Args:
            bit_string (str): 二进制序列

        Returns:
            bytes: 生成的字节序列
        """
        padding = 8 - len(bit_string) % 8
        bit_string += '0' * padding
        byte_data = bytearray()
        for i in range(0, len(bit_string), 8):
            byte_data.append(int(bit_string[i:i + 8], 2))
        return bytes([padding]) + byte_data

    def compress(self) -> None:
        if not self.input_path.exists():
            print(f"输入文件不存在：{self.input_path}")
            return

        with self.input_path.open('r', encoding='utf-8') as file:
            content = file.read()

        freq_dict = self._get_frequency(content)
        self.tree.build_tree(freq_dict)
        codes = self.tree.get_codes()

        encoded_content = ''.join(codes[char] for char in content)
        byte_data = self._to_byte_array(encoded_content)

        with self.output_path.open('wb') as file:
            json_metadata = json.dumps(freq_dict).encode('utf-8')
            file.write(len(json_metadata).to_bytes(4, 'big'))
            file.write(json_metadata)
            file.write(byte_data)

        print(f"压缩完成，保存至：{self.output_path}")

if __name__ == "__main__":
    input_path = Path("HP.txt")
    output_path = Path("compressed_HP.txt")
    compressor = HuffmanCompressor(input_path, output_path)
    compressor.compress()
